# memory_systems/experience_metadata_db.py
import sqlite3
from loguru import logger
from datetime import datetime

class ExperienceMetadataDB:
    def __init__(self, db_path="memory_systems/experience_metadata.db"):
        self.db_path = db_path
        self._create_table()

    def _create_table(self):
        conn = None
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS experiences (
                    experience_id TEXT PRIMARY KEY,
                    success_count INTEGER DEFAULT 0,
                    failure_count INTEGER DEFAULT 0,
                    total_references INTEGER DEFAULT 0,
                    average_cost REAL DEFAULT 0.0,
                    average_duration REAL DEFAULT 0.0,
                    last_updated TEXT
                )
            """)
            conn.commit()
            logger.info(f"Database table 'experiences' created or already exists at {self.db_path}")
        except sqlite3.Error as e:
            logger.error(f"Error creating database table: {e}")
        finally:
            if conn:
                conn.close()

    def upsert_experience_metadata(self, experience_id: str, success: bool, cost: float, duration: float):
        conn = None
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            current_time = datetime.now().isoformat()

            # Check if the experience_id already exists
            cursor.execute("SELECT success_count, failure_count, total_references, average_cost, average_duration FROM experiences WHERE experience_id = ?", (experience_id,))
            row = cursor.fetchone()

            if row:
                # Update existing record
                success_count, failure_count, total_references, avg_cost, avg_duration = row
                
                new_success_count = success_count + (1 if success else 0)
                new_failure_count = failure_count + (0 if success else 1)
                new_total_references = total_references + 1
                
                # Update average cost and duration
                # Simple moving average for now, can be improved later
                new_average_cost = (avg_cost * total_references + cost) / new_total_references
                new_average_duration = (avg_duration * total_references + duration) / new_total_references

                cursor.execute("""
                    UPDATE experiences SET
                        success_count = ?,
                        failure_count = ?,
                        total_references = ?,
                        average_cost = ?,
                        average_duration = ?,
                        last_updated = ?
                    WHERE experience_id = ?
                """, (new_success_count, new_failure_count, new_total_references, new_average_cost, new_average_duration, current_time, experience_id))
                logger.info(f"Updated metadata for experience_id: {experience_id}")
            else:
                # Insert new record
                success_count = 1 if success else 0
                failure_count = 0 if success else 1
                total_references = 1
                average_cost = cost
                average_duration = duration

                cursor.execute("""
                    INSERT INTO experiences (experience_id, success_count, failure_count, total_references, average_cost, average_duration, last_updated)
                    VALUES (?, ?, ?, ?, ?, ?, ?)
                """, (experience_id, success_count, failure_count, total_references, average_cost, average_duration, current_time))
                logger.info(f"Inserted new metadata for experience_id: {experience_id}")
            
            conn.commit()
        except sqlite3.Error as e:
            logger.error(f"Error upserting experience metadata for {experience_id}: {e}")
        finally:
            if conn:
                conn.close()

    def get_experience_metadata(self, experience_id: str) -> dict | None:
        conn = None
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            cursor.execute("SELECT * FROM experiences WHERE experience_id = ?", (experience_id,))
            row = cursor.fetchone()
            if row:
                columns = [description[0] for description in cursor.description]
                return dict(zip(columns, row))
            return None
        except sqlite3.Error as e:
            logger.error(f"Error retrieving experience metadata for {experience_id}: {e}")
            return None
        finally:
            if conn:
                conn.close()

    def get_all_experience_metadata(self) -> list[dict]:
        conn = None
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            cursor.execute("SELECT * FROM experiences")
            rows = cursor.fetchall()
            columns = [description[0] for description in cursor.description]
            return [dict(zip(columns, row)) for row in rows]
        except sqlite3.Error as e:
            logger.error(f"Error retrieving all experience metadata: {e}")
            return []
        finally:
            if conn:
                conn.close()
